#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
方诺双臂VR远程控制脚本 - 真实机械臂版本
基于FangnuoCar类实现，支持真实机械臂的VR远程控制

增加了数据记录功能，用于分析逆运动学计算耗时
"""
import traceback

import numpy as np
import cv2
import time
import logging
import math
import os
import sys
import argparse
import threading
import signal
from multiprocessing import shared_memory, Array, Lock, Value
from utils.logger import run_time

# 导入键盘处理模块
from utils.keyboard_handler import KeyboardHandler

# 添加父目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

# 添加base_controller目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'base_controller'))

# 导入自定义模块
sys.path.append(current_dir)
from utils.ik_data_recorder import IKDataRecorder

# 导入必要的模块
from teleop.utils.logger import get_logger
from teleop.open_television.tv_wrapper import TeleVisionWrapper
from teleop.image_server.image_client import ImageClient
from teleop.utils.episode_writer import EpisodeWriter
from teleop.robot_control.fangnuo_car import FangnuoCar, Gripper_Controller
from teleop.robot_control.fangnuo_car_ik_new import FangNuoCarIK
from teleop.utils.mat_tool import print_pose_info

# 导入底盘控制器
try:
    from serial_controller import Controller
except ImportError as e:
    print(f"警告: 无法导入底盘控制器 - {str(e)}")
    print("底盘控制功能将不可用")
    Controller = None

# 全局变量，用于信号处理
global_ik_data_recorder = None

# 信号处理函数
def signal_handler(sig, frame):
    logger = get_logger('FangnuoTeleopReal')
    logger.warning('检测到信号中断，尝试保存数据...')
    
    # 保存IK数据
    if global_ik_data_recorder is not None:
        try:
            saved_file = global_ik_data_recorder.save()
            logger.info(f'IK数据已保存到: {saved_file}')
            logger.info(f'可以使用以下命令分析IK数据: python teleop/ik_replay_analyzer.py --data-file {saved_file}')
        except Exception as e:
            logger.error(f'保存IK数据时出错: {e}')
    
    # 退出程序
    sys.exit(0)

def main():
    # 获取日志记录器
    logger = get_logger('FangnuoTeleopReal')
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='FangNuo Teleop Real Arm')
    parser.add_argument('--task-dir', type=str, default='./dataset/task_data', help='任务数据目录')
    parser.add_argument('--frequency', type=float, default=30.0, help='控制频率')
    parser.add_argument('--record', action='store_true', default=False, help='是否记录数据')
    parser.add_argument('--record-ik', action='store_true', default=False, help='是否记录逆运动学数据')
    parser.add_argument('--hand', type=str, default='gripper', choices=['gripper', 'hand'], help='末端执行器类型')
    parser.add_argument('--sim', action='store_true', default=False, help='是否使用仿真模式')
    parser.add_argument('--server-address', type=str, default='192.168.1.3', help='图像服务器地址')
    parser.add_argument('--server-port', type=int, default=5555, help='图像服务器端口')
    args = parser.parse_args()
    logger.info(f"参数: {args}\n")
    
    # 确定是否使用真实机械臂
    use_real_arms = not args.sim
    logger.info(f"机械臂模式: {'真实机械臂' if use_real_arms else '模拟机械臂'}")

    # 图像客户端配置
    img_config = {
        'fps': 30,
        'head_camera_type': 'opencv',
        'head_camera_image_shape': [480, 640],  # 头部相机分辨率
        'head_camera_id_numbers': [4],
        'wrist_camera_type': 'opencv',
        'wrist_camera_image_shape': [480, 640],  # 手腕相机分辨率
        'wrist_camera_id_numbers': [6, 8],
    }
    
    # 判断是否为双目相机
    ASPECT_RATIO_THRESHOLD = 2.0  # 如果宽高比超过此值，则认为是双目相机
    BINOCULAR = False
    
    # 判断是否有手腕相机
    if 'wrist_camera_type' in img_config:
        WRIST = True
    else:
        WRIST = False
    
    # 设置图像形状
    if BINOCULAR and not (img_config['head_camera_image_shape'][1] / img_config['head_camera_image_shape'][0] > ASPECT_RATIO_THRESHOLD):
        tv_img_shape = (img_config['head_camera_image_shape'][0], img_config['head_camera_image_shape'][1] * 2, 3)
    else:
        tv_img_shape = (img_config['head_camera_image_shape'][0], img_config['head_camera_image_shape'][1], 3)

    # 创建共享内存
    tv_img_shm = shared_memory.SharedMemory(create=True, size=np.prod(tv_img_shape) * np.uint8().itemsize)
    tv_img_array = np.ndarray(tv_img_shape, dtype=np.uint8, buffer=tv_img_shm.buf)
    
    # 创建时间戳共享内存，用于图像同步
    timestamp_shm = shared_memory.SharedMemory(create=True, size=8)  # 8字节用于存储float64时间戳
    timestamp_array = np.ndarray((1,), dtype=np.float64, buffer=timestamp_shm.buf)
    timestamp_array[0] = 0.0  # 初始化时间戳

    # 初始化图像客户端
    if WRIST:
        wrist_img_shape = (img_config['wrist_camera_image_shape'][0], img_config['wrist_camera_image_shape'][1] * 2, 3)
        wrist_img_shm = shared_memory.SharedMemory(create=True, size=np.prod(wrist_img_shape) * np.uint8().itemsize)
        wrist_img_array = np.ndarray(wrist_img_shape, dtype=np.uint8, buffer=wrist_img_shm.buf)
        img_client = ImageClient(tv_img_shape=tv_img_shape, tv_img_shm_name=tv_img_shm.name, 
                                wrist_img_shape=wrist_img_shape, wrist_img_shm_name=wrist_img_shm.name,
                                server_address=args.server_address, port=args.server_port, Unit_Test=False, image_show=False,
                                timestamp_shm_name=timestamp_shm.name)
        logger.info(f"[图像客户端] 连接到服务器: {args.server_address}:{args.server_port}")
    else:
        img_client = ImageClient(tv_img_shape=tv_img_shape, tv_img_shm_name=tv_img_shm.name,
                               server_address=args.server_address, port=args.server_port,
                               timestamp_shm_name=timestamp_shm.name)
        logger.info(f"[图像客户端] 连接到服务器: {args.server_address}:{args.server_port}")

    # 启动图像接收线程
    image_receive_thread = threading.Thread(target=img_client.receive_process, daemon=True)
    image_receive_thread.daemon = True
    image_receive_thread.start()

    # 初始化VR设备接口
    tv_wrapper = TeleVisionWrapper(BINOCULAR, tv_img_shape, tv_img_shm.name)
    
    # 创建共享内存，用于存储right_squeeze_value
    right_squeeze_value_shm = Value('d', 0.0, lock=True)  # 使用双精度浮点数存储right_squeeze_value

    # 初始化机械臂控制器
    logger.info(f"初始化机械臂控制器...")
    arm_ctrl = FangnuoCar(use_real_arms=use_real_arms)
    arm_ik = FangNuoCarIK(Unit_Test=False, Visualization=False)
    
    # 初始化底盘控制器
    # base_controller = None
    # if Controller is not None:
    #     try:
    #         logger.info("初始化底盘控制器...")
    #         base_controller = Controller(
    #             serial_port='/dev/ttyUSB0',  # 根据实际情况修改串口设备
    #             baud_rate=38400,
    #             socket_host='0.0.0.0',
    #             socket_port=8888
    #         )
    #         base_controller.start()
    #         logger.info("底盘控制器初始化完成")
            
    #         # 将base_controller添加到tv_wrapper中，以便在回调中使用
    #         self.base_controller = base_controller
    #         logger.info("底盘控制器已添加到tv_wrapper")
    #     except Exception as e:
    #         logger.error(f"初始化底盘控制器失败: {str(e)}")
    # else:
    #     logger.warning("未找到底盘控制器模块，底盘控制功能将不可用")
    
    # 初始化IK数据记录器（如果启用）
    ik_data_recorder = None
    if args.record_ik:
        logger.info("初始化IK数据记录器...")
        # 使用自动保存功能，每50条数据自动保存一次
        ik_data_recorder = IKDataRecorder(auto_save_interval=50)
        global global_ik_data_recorder
        global_ik_data_recorder = ik_data_recorder
        ik_data_session = ik_data_recorder.start_new_session()
        logger.info(f"IK数据将保存到: {ik_data_session}")
        
        # 注册信号处理函数，确保在程序被中断时保存数据
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
    else:
        logger.info("IK数据记录功能未启用，使用--record_ik参数启用")


    # 初始化手部控制器
    if args.hand == "gripper":
        left_hand_array = Array('d', 75, lock=True)
        right_hand_array = Array('d', 75, lock=True)
        dual_gripper_data_lock = Lock()
        dual_gripper_state_array = Array('d', 2, lock=False)   # 当前左右夹爪状态数据
        dual_gripper_action_array = Array('d', 2, lock=False)  # 当前左右夹爪动作数据
        
    # 初始化数据记录器
    recorder = None
    if args.record:
        recorder = EpisodeWriter(task_dir=args.task_dir, frequency=args.frequency, rerun_log=True)
    
    # 初始化键盘处理器
    keyboard_handler = KeyboardHandler(arm_ctrl=arm_ctrl, recorder=recorder, record_enabled=args.record)
    
    # 记录上一次的时间戳，用于检测新图像
    last_timestamp = 0.0
    
    try:
        # 等待用户输入启动信号
        user_input = input("请输入启动信号（输入'r'启动程序）:\n")
        if user_input.lower() == 'r':
            # 设置机械臂速度逐渐增加到最大值
            if not use_real_arms:  
                arm_ctrl.speed_gradual_max()

            ACTIVATE_HAND = False
            
            # 主循环
            while keyboard_handler.running:
                start_time = time.time()
                
                # 检查是否有新图像
                current_timestamp = timestamp_array[0]
                has_new_image = current_timestamp > last_timestamp
                if has_new_image:
                    last_timestamp = current_timestamp
                    # 处理新图像
                    logger.debug(f"[图像同步] 接收到新图像，时间戳: {current_timestamp:.6f}")
                
                # 获取手柄数据和按钮状态
                head_rmat, left_wrist, right_wrist, left_hand, right_hand, controller_data = tv_wrapper.get_data()
                
                # 从手柄数据中获取按钮状态
                right_squeeze_value = controller_data.get('right_squeeze_value', 0.0)
                # 获取其他按钮状态
                right_a_button = controller_data.get('right_a_button', False)
                right_b_button = controller_data.get('right_b_button', False)
                right_thumbstick_x = controller_data.get('right_thumbstick_x', 0.0) # 遥控底盘左右 [1,-1]
                right_thumbstick_y = controller_data.get('right_thumbstick_y', 0.0) # 遥控底盘前后 [1,-1]
                
                # 控制底盘移动
                # 设置最大线速度和角速度
                MAX_LINEAR_VEL = 0.7  # 最大线速度 (m/s)
                MAX_ANGULAR_VEL = 0.5  # 最大角速度 (rad/s)
                
                # 计算线速度 (vx) 和角速度 (vw)
                # 前后摇杆控制线速度 (反向，因为摇杆向上为负，向下为正)
                vx = -right_thumbstick_y * MAX_LINEAR_VEL
                # 左右摇杆控制角速度
                vw = -right_thumbstick_x * MAX_ANGULAR_VEL
                
                # 调用底盘控制命令 (vy固定为0，因为不需要横向移动)
                # if hasattr(self, 'base_controller') and self.base_controller is not None:
                #     try:
                #         self.base_controller.handle_command(vx, 0.0, vw)
                #     except Exception as e:
                #         logger.error(f"发送底盘控制命令失败: {str(e)}")
                
                # 获取左手柄按钮状态
                left_squeeze = controller_data.get('left_squeeze', False)
                left_squeeze_value = controller_data.get('left_squeeze_value', 0.0)
                left_x_button = controller_data.get('left_x_button', False)
                # left_y_button = controller_data.get('left_y_button', False)
                
                MAX_GRIPPER_OPEN = 70.0  # 夹爪最大开度，单位mm

                # 将右手握力值转换为夹爪位置（反转逻辑：默认打开，握力值增加时关闭）
                right_gripper_position = MAX_GRIPPER_OPEN * (1.0 - right_squeeze_value)
                # 设置夹爪力矩，握力值越大，力矩越大
                right_gripper_force = 500 + int(right_squeeze_value * 2000)  # 力矩范围500-2500
                # 控制右臂夹爪
                arm_ctrl.control_right_gripper(right_gripper_position, right_gripper_force)
                logger.info(f"右手夹爪控制: 握力值={right_squeeze_value:.2f}, 夹爪位置={right_gripper_position:.2f}mm, 力矩={right_gripper_force}")

     
                # 将左手握力值转换为夹爪位置（反转逻辑：默认打开，握力值增加时关闭）
                left_gripper_position = MAX_GRIPPER_OPEN * (1.0 - left_squeeze_value)
                # 设置夹爪力矩，握力值越大，力矩越大
                left_gripper_force = 500 + int(left_squeeze_value * 2000)  # 力矩范围500-2500
                # 控制左臂夹爪
                arm_ctrl.control_left_gripper(left_gripper_position, left_gripper_force)
                logger.info(f"左手夹爪控制: 握力值={left_squeeze_value:.2f}, 夹爪位置={left_gripper_position:.2f}mm, 力矩={left_gripper_force}")
            
                # 获取当前夹爪状态（仅用于日志记录）
                gripper_state = arm_ctrl.get_gripper_state()
                logger.debug(f"当前夹爪状态: 左={gripper_state['left_gripper']:.2f}mm, 右={gripper_state['right_gripper']:.2f}mm, \
                    joystick:left_squeeze={left_squeeze_value:.2f}, right_squeeze={right_squeeze_value:.2f} \
                        left_x_gohome={left_x_button}, right_a_activate={right_a_button}")
                
                if right_a_button:
                    ACTIVATE_HAND = True
                    logger.info(f"检测到右手柄B按钮，切换手部控制模式: {'激活手部' if ACTIVATE_HAND else '不激活手部'}")
                    arm_ctrl.enable_dual_arm()

                if right_b_button:
                    ACTIVATE_HAND = False
                    logger.info(f"检测到右手柄B按钮，切换手部控制模式: {'激活手部' if ACTIVATE_HAND else '不激活手部'}")
                    
                
                # 检查是否按下了重置按钮（使用A按钮作为重置按钮）
                if left_x_button:
                    logger.info("检测到左手柄A按钮，重置机械臂到初始位置")
                    ACTIVATE_HAND = False
                    arm_ctrl.go_home()
                    time.sleep(5.0)
                    arm_ctrl.stop()
                    time.sleep(1.0)  # 给重置命令一些时间生效
                
                # 更新共享内存中的right_squeeze_value
                with right_squeeze_value_shm.get_lock():
                    right_squeeze_value_shm.value = right_squeeze_value
                
                logger.info(f"-----left_wrist: {print_pose_info(left_wrist)} left_squeeze_value: {left_squeeze_value}")
                logger.info(f"-----right_wrist: {print_pose_info(right_wrist)} right_squeeze_value: {right_squeeze_value}")
          
                
                # 显示夹爪状态
                gripper_state = arm_ctrl.get_gripper_state()
                logger.info(f"-----夹爪状态: 左={gripper_state['left_gripper']:.2f}mm, 右={gripper_state['right_gripper']:.2f}mm")

                # 发送手部骨骼数据到手部控制器
                if args.hand:
                    left_hand_array[:] = left_hand.flatten()
                    right_hand_array[:] = right_hand.flatten()

                # 获取当前机械臂状态
                arm_status = arm_ctrl.status()
                
                # 确保我们获取完整的12个关节数据（左右臂各6个关节）
                left_arm_q = arm_status['left_arm']['position']
                right_arm_q = arm_status['right_arm']['position']
                left_arm_dq = arm_status['left_arm']['velocity']
                right_arm_dq = arm_status['right_arm']['velocity']
                
                # 检查并确保数据完整性
                if len(left_arm_q) != 6 or len(right_arm_q) != 6 or len(left_arm_dq) != 6 or len(right_arm_dq) != 6:
                    logger.warning(f"关节数据不完整: 左臂位置={len(left_arm_q)}, 右臂位置={len(right_arm_q)}, 左臂速度={len(left_arm_dq)}, 右臂速度={len(right_arm_dq)}")
                    # 如果数据不完整，使用零填充到正确长度
                    left_arm_q = np.pad(left_arm_q, (0, max(0, 6-len(left_arm_q))), 'constant')
                    right_arm_q = np.pad(right_arm_q, (0, max(0, 6-len(right_arm_q))), 'constant')
                    left_arm_dq = np.pad(left_arm_dq, (0, max(0, 6-len(left_arm_dq))), 'constant')
                    right_arm_dq = np.pad(right_arm_dq, (0, max(0, 6-len(right_arm_dq))), 'constant')
                
                # 合并左右臂数据为12个关节的完整数据
                current_lr_arm_q = np.concatenate((left_arm_q, right_arm_q))
                current_lr_arm_dq = np.concatenate((left_arm_dq, right_arm_dq))
                
                if ACTIVATE_HAND:
         

                    # 使用电机数据和手腕姿态求解逆运动学，然后使用结果控制机械臂
                    time_ik_start = time.time()
                    sol_q, sol_tauff = arm_ik.solve_ik(left_wrist, right_wrist, current_lr_arm_q, current_lr_arm_dq)
                    time_ik_end = time.time()
                    ik_time = time_ik_end - time_ik_start
                    logger.info(f"ik计算时间:\t{round(ik_time, 6)} 更新机器人:{sol_q}")
                    # display end frame
                    arm_ik.get_frame_pose(sol_q)

                    if sol_q is None:
                        continue

                    # 记录IK数据（如果启用）
                    if args.record_ik and ik_data_recorder is not None:
                        ik_data_recorder.record(left_wrist, right_wrist, current_lr_arm_q, current_lr_arm_dq, sol_q)
                    
                    # 控制机械臂
                    # arm_ctrl.ctrl_dual_arm(sol_q, sol_tauff)

                    # 记录训练数据集
                    if recorder and keyboard_handler.is_recording():  # 只在recording为True时记录
                        # 准备图像数据
                        colors = {}
                        depths = {}
                        
                        # 获取当前图像
                        current_tv_image = np.copy(tv_img_array)
                        
                        # 根据BINOCULAR和WRIST配置添加相应的图像
                        if BINOCULAR:
                            colors[f"color_0"] = current_tv_image[:, :tv_img_shape[1]//2]  # 左侧头部相机
                            colors[f"color_1"] = current_tv_image[:, tv_img_shape[1]//2:]  # 右侧头部相机
                        else:
                            colors[f"color_0"] = current_tv_image
                            
                        if WRIST:
                            current_wrist_image = np.copy(wrist_img_array)
                            colors[f"color_2"] = current_wrist_image[:, :wrist_img_shape[1]//2]  # 左腕相机
                            colors[f"color_3"] = current_wrist_image[:, wrist_img_shape[1]//2:]  # 右腕相机
                        
                        # 准备状态数据
                        states = {
                            "left_arm": {                                                                    
                                "qpos":   left_arm_q.tolist(),    # numpy.array -> list
                                "qvel":   left_arm_dq.tolist(),                          
                                "torque": [],                        
                            }, 
                            "right_arm": {                                                                    
                                "qpos":   right_arm_q.tolist(),       
                                "qvel":   right_arm_dq.tolist(),                          
                                "torque": [],                         
                            },                        
                            "left_hand": {                                                                    
                                "qpos":   [gripper_state['left_gripper']],           
                                "qvel":   [],                           
                                "torque": [],                          
                            }, 
                            "right_hand": {                                                                    
                                "qpos":   [gripper_state['right_gripper']],       
                                "qvel":   [],                           
                                "torque": [],  
                            }, 
                            "body": None, 
                        }
                        
                        # 准备动作数据
                        actions = {
                            "left_arm": {                                   
                                "qpos":   sol_q[:6].tolist() if sol_q is not None else [],       
                                "qvel":   [],       
                                "torque": sol_tauff[:6].tolist() if sol_tauff is not None else [],      
                            }, 
                            "right_arm": {                                   
                                "qpos":   sol_q[6:].tolist() if sol_q is not None else [],       
                                "qvel":   [],       
                                "torque": sol_tauff[6:].tolist() if sol_tauff is not None else [],       
                            },                         
                            "left_hand": {                                   
                                "qpos":   [left_gripper_position],       
                                "qvel":   [],       
                                "torque": [left_gripper_force],       
                            }, 
                            "right_hand": {                                   
                                "qpos":   [right_gripper_position],       
                                "qvel":   [],       
                                "torque": [right_gripper_force], 
                            }, 
                            "body": None, 
                        }
                        
                        # 调用recorder.add_item记录数据
                        recorder.add_item(colors=colors, depths=depths, states=states, actions=actions)

                else:
                    logger.debug("右手柄侧握键未被按下")

                # 准备显示图像
                tv_resized_image = cv2.resize(tv_img_array, (tv_img_shape[1] // 2, tv_img_shape[0] // 2))
                
                # 处理键盘输入、显示状态信息并显示图像
                keyboard_handler.process_keyboard_input(tv_resized_image)
    

                # 控制循环频率
                current_time = time.time()
                time_elapsed = current_time - start_time
                sleep_time = max(0, (1 / float(args.frequency)) - time_elapsed)
                time.sleep(sleep_time)
                logger.info(f"time_elapsed: {time_elapsed}")
                
    except KeyboardInterrupt:
        logger.warning("检测到键盘中断，退出程序...")
    except Exception as e:
        logger.error(f"发生异常: {e}")
        logger.error("详细堆栈信息:\n%s", traceback.format_exc())
    finally:
        # 停止机械臂
        logger.info("程序结束，停止机械臂...")
        try:
            arm_ctrl.go_home()
            time.sleep(5.0)
            arm_ctrl.stop()
            
            # 保存IK数据（如果启用）
            if args.record_ik and ik_data_recorder is not None:
                logger.info("保存IK数据...")
                saved_file = ik_data_recorder.save_session()
                logger.info(f"IK数据已保存到: {saved_file}")
                logger.info(f"可以使用以下命令进行数据分析: python teleop/ik_replay_analyzer.py --data-file={saved_file}")

        except:
            pass
        
        # 清理资源
        logger.info("清理资源...")
        cv2.destroyAllWindows()
        
        # 释放共享内存
        try:
            tv_img_shm.close()
            tv_img_shm.unlink()
            timestamp_shm.close()
            timestamp_shm.unlink()
            if WRIST:
                wrist_img_shm.close()
                wrist_img_shm.unlink()
        except:
            pass
        
        logger.info("程序已安全退出")

if __name__ == '__main__':
    main()
